# Copyright (c) 2024, NVIDIA CORPORATION.
import io
import os
from collections.abc import Mapping
from enum import IntEnum
from typing import Any, Literal, Self, TypeAlias, overload

from pylibcudf.column import Column
from pylibcudf.io.datasource import Datasource
from pylibcudf.table import Table

class JSONRecoveryMode(IntEnum):
    FAIL = ...
    RECOVER_WITH_NULL = ...

class CompressionType(IntEnum):
    NONE = ...
    AUTO = ...
    SNAPPY = ...
    GZIP = ...
    BZIP2 = ...
    BROTLI = ...
    ZIP = ...
    XZ = ...
    ZLIB = ...
    LZ4 = ...
    LZO = ...
    ZSTD = ...

class ColumnEncoding(IntEnum):
    USE_DEFAULT = ...
    DICTIONARY = ...
    PLAIN = ...
    DELTA_BINARY_PACKED = ...
    DELTA_LENGTH_BYTE_ARRAY = ...
    DELTA_BYTE_ARRAY = ...
    BYTE_STREAM_SPLIT = ...
    DIRECT = ...
    DIRECT_V2 = ...
    DICTIONARY_V2 = ...

class DictionaryPolicy(IntEnum):
    NEVER = ...
    ADAPTIVE = ...
    ALWAYS = ...

class StatisticsFreq(IntEnum):
    STATISTICS_NONE = ...
    STATISTICS_ROWGROUP = ...
    STATISTICS_PAGE = ...
    STATISTICS_COLUMN = ...

class QuoteStyle(IntEnum):
    MINIMAL = ...
    ALL = ...
    NONNUMERIC = ...
    NONE = ...

ColumnNameSpec: TypeAlias = tuple[str, list[ColumnNameSpec]]
ChildNameSpec: TypeAlias = Mapping[str, ChildNameSpec]

class PartitionInfo:
    def __init__(self, start_row: int, num_rows: int): ...

class TableInputMetadata:
    def __init__(self, table: Table): ...
    @property
    def column_metadata(self) -> list[ColumnInMetadata]: ...

class ColumnInMetadata:
    def set_name(self, name: str) -> Self: ...
    def set_nullability(self, nullable: bool) -> Self: ...
    def set_list_column_as_map(self) -> Self: ...
    def set_int96_timestamps(self, req: bool) -> Self: ...
    def set_decimal_precision(self, precision: int) -> Self: ...
    def child(self, i: int) -> Self: ...
    def set_output_as_binary(self, binary: bool) -> Self: ...
    def set_type_length(self, type_length: int) -> Self: ...
    def set_skip_compression(self, skip: bool) -> Self: ...
    def set_encoding(self, encoding: ColumnEncoding) -> Self: ...
    def get_name(self) -> str: ...

class TableWithMetadata:
    tbl: Table
    def __init__(
        self, tbl: Table, column_names: list[ColumnNameSpec]
    ) -> None: ...
    @property
    def columns(self) -> list[Column]: ...
    @overload
    def column_names(self, include_children: Literal[False]) -> list[str]: ...
    @overload
    def column_names(
        self, include_children: Literal[True]
    ) -> list[ColumnNameSpec]: ...
    @overload
    def column_names(
        self, include_children: bool = False
    ) -> list[str] | list[ColumnNameSpec]: ...
    @property
    def child_names(self) -> ChildNameSpec: ...
    @property
    def per_file_user_data(self) -> list[Mapping[str, str]]: ...

class SourceInfo:
    def __init__(
        self, sources: list[str] | list[os.PathLike[Any]] | list[Datasource]
    ) -> None: ...

class SinkInfo:
    def __init__(
        self,
        sinks: list[os.PathLike[Any]]
        | list[io.StringIO]
        | list[io.BytesIO]
        | list[io.TextIOBase]
        | list[str],
    ) -> None: ...
